import torch
import torch.nn as nn
from encoder import ResNetUNet
from decoder import Decoder
from faceswap_pytorch.fs_models import Autoencoder
from SimSwap_main.simswap_model import Simswapencoder

#模型主功能类
class FaceTagger(nn.Module) :
    def __init__(self, message_size, device, 
                 GAN_sim="faceswap", decoder="ResnetUnet"):
        
        super().__init__()
        self.encoder = ResNetUNet(message_size)
        #df_model: nn.Moudle
        self.decoder = Decoder(message_size)
        self.faceswap_PATH = './faceswap_pytorch/checkpoint/autoencoder.t7'

        #加载换脸模型
        try:
            #模型文件
            if GAN_sim == "faceswap":
                df_path = self.faceswap_PATH
                checkpoint = torch.load(df_path)
                self.df_model = Autoencoder()
                self.df_model.load_state_dict(checkpoint['state'])
            else:
                #self.df_model = create_model(TestOptions().parse())
                self.df_model = Simswapencoder()
        except: 
            print("model load fail")

        #加载模型至对应设备
        if device :
            self.encoder = self.encoder.to(device)
            self.df_model = self.df_model.to(device)
            self.decoder = self.decoder.to(device)
        # for param in self.df_model.features.parameter() :
        #     param.requires_grad = False
        # self.df_model.eval()
        
        
        nets = self.df_model
        if not isinstance(nets, list):
            nets = [nets]
        for net in nets:
            if net is not None:
                for param in net.parameters():
                    param.requires_grad = False
    
    #引用编码--换脸--解码 方法
    def encode(self, x, message) :
        return self.encoder(x, message)
        
    def faceswap(self, x, type) : 
        return self.df_model(x, type)

    def simswap(self, id, spc):
        return self.df_model.inference(id, spc, "ATT")
    
    def decode(self, x) :
        # print(x.shape)
        return self.decoder(x)